package com.clei.datastructure;

import com.clei.utils.CollectionUtil;
import com.clei.utils.PrintUtil;
import com.clei.utils.StringUtil;
import com.clei.utils.container.Pair;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * 前缀树
 * 单词查找树、字典树
 * 字符组词 适合中文
 *
 * @author Y
 * @date 2024-10-22
 **/
public class TrieTreeNode {

    /**
     * 字符
     */
    private char c;

    /**
     * 词结束
     */
    private boolean endOfWord;

    /**
     * 子结点
     */
    private Map<Character, TrieTreeNode> children;

    public static void main(String[] args) {
        String content = "烛龙栖寒门，光曜犹旦开\n" +
                "日月照之何不及此？唯有北风号怒天上来\n" +
                "燕山雪花大如席，片片吹落轩辕台\n" +
                "幽州思妇十二月，停歌罢笑双蛾摧\n" +
                "倚门望行人，念君长城苦寒良可哀\n" +
                "别时提剑救边去，遗此虎纹金鞞靫\n" +
                "中有一双白羽箭，蜘蛛结网生尘埃\n" +
                "箭空在，人今战死不复回\n" +
                "不忍见此物，焚之已成灰\n" +
                "黄河捧土犹可塞，北风雨雪恨难裁";
        String[] wordArr = {"烛龙", "人", "日月", "北风", "北风雨雪", "轩辕台", "十二月", "行人", "金鞞靫", "战死", "不忍", "恨难裁定", "难裁"};
        // 创建树
        TrieTreeNode tree = new TrieTreeNode(Arrays.stream(wordArr).collect(Collectors.toList()));
        // 查找
        List<Pair<Integer, Integer>> matchPair = tree.search(content);
        // 展示
        display(content, matchPair);
    }

    /**
     * 展示
     *
     * @param content   content
     * @param matchPair 匹配字符
     */
    public static void display(String content, List<Pair<Integer, Integer>> matchPair) {
        if (CollectionUtil.isEmpty(matchPair)) {
            PrintUtil.print(content);
        }
        // 位置排序
        matchPair.sort(Comparator.comparingInt(Pair::getFirst));
        // append
        StringBuilder sb = new StringBuilder(content.length() + matchPair.size() * 2);
        int before = 0;
        for (Pair<Integer, Integer> pair : matchPair) {
            sb.append(content, before, pair.getFirst())
                    .append('[')
                    .append(content, pair.getFirst(), pair.getFirst() + pair.getSecond())
                    .append(']');
            before = pair.getFirst() + pair.getSecond();
        }
        sb.append(content.substring(before));
        PrintUtil.print(sb.toString());
    }

    /**
     * 添加词
     *
     * @param word 词语
     */
    public void insert(String word) {
        if (StringUtil.isEmpty(word)) {
            return;
        }
        char[] arr = word.toCharArray();
        int length = arr.length;
        int count = 1;
        Map<Character, TrieTreeNode> children = this.children;
        for (char c : arr) {
            TrieTreeNode node = children.get(c);
            if (null == node) {
                node = new TrieTreeNode(c, false, new HashMap<>());
                children.put(c, node);
            }
            // 词语
            if (count == length) {
                node.setEndOfWord(true);
            }
            children = node.getChildren();
            count++;
        }
    }

    /**
     * search
     *
     * @param content content
     * @return Pair list key->startIndex value->length
     */
    public List<Pair<Integer, Integer>> search(String content) {
        if (StringUtil.isBlank(content)) {
            return Collections.emptyList();
        }
        return search(content.toCharArray(), 0);
    }

    /**
     * search
     *
     * @param start 起始位置
     * @param arr   arr
     * @return Pair list key->startIndex value->length
     */
    public List<Pair<Integer, Integer>> search(char[] arr, int start) {
        List<Pair<Integer, Integer>> resList = new ArrayList<>();
        int arrLength = arr.length;
        Map<Character, TrieTreeNode> children = this.children;
        int matchStart = -1, matchLength = 0;
        for (int i = start; i < arrLength; i++) {
            char c = arr[i];
            TrieTreeNode node = children.get(c);
            if (null == node) {
                // 上个匹配结束
                if (matchLength > 0) {
                    resList.add(Pair.of(matchStart, matchLength));
                    // 跳转到匹配词下个字符开始匹配
                    i = matchStart + matchLength - 1;
                    children = this.children;
                    matchStart = -1;
                    matchLength = 0;
                    continue;
                }
                if (matchStart > 0) {
                    // 跳转到匹配字符下个字符开始匹配
                    i = matchStart;
                    children = this.children;
                    matchStart = -1;
                    matchLength = 0;
                    continue;
                }
                continue;
            }
            // 匹配
            // 首次匹配
            if (-1 == matchStart) {
                matchStart = i;
            }
            // 树往下走
            children = node.getChildren();
            // 词语匹配成功 继续往下走 优先匹配长词语
            if (node.isEndOfWord()) {
                matchLength = i - matchStart + 1;
            }
        }
        // 匹配到最后
        if (matchLength > 0) {
            resList.add(Pair.of(matchStart, matchLength));
            matchStart = matchStart + matchLength - 1;
        }
        // 匹配到最后还在匹配 避免错过最后的短词
        if (matchStart > 0 && matchStart < (arrLength - 1)) {
            resList.addAll(search(arr, matchStart + 1));
        }
        return resList;
    }

    public TrieTreeNode(List<String> wordList) {
        this(' ', false, new HashMap<>());
        for (String word : wordList) {
            this.insert(word);
        }
    }

    public TrieTreeNode() {
    }

    public TrieTreeNode(char c, boolean endOfWord, Map<Character, TrieTreeNode> children) {
        this.c = c;
        this.endOfWord = endOfWord;
        this.children = children;
    }

    public char getC() {
        return c;
    }

    public void setC(char c) {
        this.c = c;
    }

    public boolean isEndOfWord() {
        return endOfWord;
    }

    public void setEndOfWord(boolean endOfWord) {
        this.endOfWord = endOfWord;
    }

    public Map<Character, TrieTreeNode> getChildren() {
        return children;
    }

    public void setChildren(Map<Character, TrieTreeNode> children) {
        this.children = children;
    }
}
